package main

import (
    "bufio"
    "bytes"
    "fmt"
    "io"
    "net"
    "net/url"
    "strings"
	"net/http"
    "golang.org/x/net/html"
	"log"
	"os"
)

func main() {
	if len(os.Args) != 2 {
		fmt.Println("Usage: http_proxy <port>")
		os.Exit(1)
	}
	port := os.Args[1]
	listener, err := net.Listen("tcp", ":"+port)
	if err != nil {
		fmt.Println("Error starting proxy:", err)
		os.Exit(1)
	}
	fmt.Println("Listening on port", port)

	for {
		conn, err := listener.Accept()
		if err != nil {
			fmt.Println("Error accepting connection:", err)
			continue
		}
		go handleConnection(conn)
	}
}

func handleConnection(clientConn net.Conn) {
	defer clientConn.Close()

	clientReader := bufio.NewReader(clientConn)
	req, err := http.ReadRequest(clientReader)
	if err != nil || req.Method != "GET" {
		writeErrorResponse(clientConn, "Internal Error", http.StatusInternalServerError)
		return
	}

	forwardRequest(req, clientConn)
}


func forwardRequest(req *http.Request, clientConn net.Conn) {
    host := req.Host
    targetConn, err := net.Dial("tcp", host+":80")
    if err != nil {
        writeErrorResponse(clientConn, "Internal Error", http.StatusInternalServerError)
        return
    }
    defer targetConn.Close()

    req.Header.Set("Connection", "close")
    err = req.Write(targetConn)
    if err != nil {
        writeErrorResponse(clientConn, "Internal Error", http.StatusInternalServerError)
        return
    }
	responseBuffer := new(bytes.Buffer)
    teeReader := io.TeeReader(targetConn, responseBuffer)

	io.Copy(clientConn, teeReader)

    go func() {
        parseAndPrefetchLinks(responseBuffer)
    }()
}


func writeErrorResponse(conn net.Conn, message string, statusCode int) {
    response := fmt.Sprintf("HTTP/1.1 %d %s\r\nContent-Length: %d\r\n\r\n%s",
        statusCode, http.StatusText(statusCode), len(message), message)
    conn.Write([]byte(response))
}


func parseAndPrefetchLinks(responseBuffer *bytes.Buffer) {
    doc, err := html.Parse(responseBuffer)
    if err != nil {
        fmt.Println("Error parsing HTML:", err)
        return
    }
    var extractLinks func(*html.Node)
    extractLinks = func(n *html.Node) {
        if n.Type == html.ElementNode && n.Data == "a" {
            for _, attr := range n.Attr {
                if attr.Key == "href" && strings.HasPrefix(attr.Val, "http") {
                    go prefetchDNS(attr.Val) 
                }
            }
        }
        for c := n.FirstChild; c != nil; c = c.NextSibling {
            extractLinks(c)
        }
    }
    extractLinks(doc)
}

func prefetchDNS(link string) {
    u, err := url.Parse(link)
    if err != nil {
        log.Println("Invalid URL:", link)
        return
    }

    host := u.Host
    if strings.Contains(host, ":") { // 如果有端口，分离主机名和端口
        host, _, err = net.SplitHostPort(host)
        if err != nil { 
            log.Println("Error parsing host:", err)
            return
        }
    }

    _, err = net.LookupHost(host)
    if err != nil {
        fmt.Println("Error during DNS lookup:", err)
    }
}
